# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Controller
"""

from typing import Union, Dict, List
from numpy import ndarray

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from ..system import Molecule
from ..function.functions import get_integer, get_tensor, get_arguments


class Controller(nn.Module):
    r"""
    Base class for the controller module in MindSPONGE.
    The `Controller` used in `Updater` to control the values of seven variables during the simulation
    process: coordinate, velocity, force, energy, kinetics, virial and pbc_box.

    Args:
        system(Molecule):   Simulation system
        control_step(int):  Step interval for controller execution. Default: 1

    Inputs:
        - **coordinate** (Tensor) - Tensor of shape `(B, A, D)`. Data type is float.
        - **velocity** (Tensor) - Tensor of shape `(B, A, D)`. Data type is float.
        - **force** (Tensor) - Tensor of shape `(B, A, D)`. Data type is float.
        - **energy** (Tensor) - Tensor of shape `(B, 1)`. Data type is float.
        - **kinetics** (Tensor) - Tensor of shape `(B, D)`. Data type is float.
        - **virial** (Tensor) - Tensor of shape `(B, D)`. Data type is float.
        - **pbc_box** (Tensor) - Tensor of shape `(B, D)`. Data type is float.
        - **step** (int) - Simulation step. Default: 0

    Outputs:
        - coordinate, Tensor of shape `(B, A, D)`. Data type is float.
        - velocity, Tensor of shape `(B, A, D)`. Data type is float.
        - force, Tensor of shape `(B, A, D)`. Data type is float.
        - energy, Tensor of shape `(B, 1)`. Data type is float.
        - kinetics, Tensor of shape `(B, D)`. Data type is float.
        - virial, Tensor of shape `(B, D)`. Data type is float.
        - pbc_box, Tensor of shape `(B, D)`. Data type is float.
    """
    def __init__(self,
                 system: Molecule,
                 control_step: int = 1,
                 **kwargs,
                 ):

        super().__init__()
        self._kwargs = get_arguments(locals(), kwargs)

        self.system = system
        self.num_walker = self.system.num_walker
        self.num_atoms = system.num_atoms
        self.dimension = system.dimension

        self.sys_dofs = system.degrees_of_freedom
        self.degrees_of_freedom = system.degrees_of_freedom

        self.time_step = get_tensor(1e-3, dtype=torch.float32)

        self._coordinate = self.system.coordinate
        self._pbc_box = self.system.pbc_box

        self.units = self.system.units
        self.kinetic_unit_scale = get_tensor(self.units.kinetic_ref, dtype=torch.float32)
        self.press_unit_scale = get_tensor(self.units.pressure_ref, dtype=torch.float32)

        # (B, A)
        self.atom_mass = self.system.atom_mass
        self.inv_mass = self.system.inv_mass
        # (B, A, 1)
        self._atom_mass = self.atom_mass.unsqueeze(-1)
        self._inv_mass = self.inv_mass.unsqueeze(-1)

        # (B, 1)
        self.system_mass = self.system.system_mass
        self.system_natom = self.system.system_natom

        self.control_step = get_integer(control_step)
        if self.control_step <= 0:
            raise ValueError('The "control_step" must be larger than 0!')

        self.num_constraints = 0

    @property
    def boltzmann(self) -> float:
        """
        Boltzmann constant in current unit.

        Returns:
            float, Boltzmann constant in current unit.
        """
        return self.units.boltzmann

    def set_time_step(self, dt: float):
        r"""
        Set simulation time step.

        Args:
            dt(float):  Time step.
        """
        self.time_step = get_tensor(dt, torch.float32)
        return self

    def set_degrees_of_freedom(self, dofs: int):
        """
        Set degrees of freedom (DOFs).

        Args:
            dofs(int):  Degrees of freedom.
        """
        self.degrees_of_freedom = get_integer(dofs)
        return self

    def update_coordinate(self, coordinate: Tensor) -> Tensor:
        r"""
        Update the coordinate of the simulation system.

        Args:
            coordinate(Tensor): Tensor of atomic coordinates. Tensor shape is `(B, A, D)`.
                                Data type is float.

        Returns:
            Tensor, has the same data type and shape as original `coordinate`.
        """
        self._coordinate.copy_(coordinate)
        return self._coordinate

    def update_pbc_box(self, pbc_box: Tensor) -> Tensor:
        r"""
        Update the parameter of PBC box.

        Args:
            pbc_box(Tensor):    Tensor of PBC box. Tensor shape is `(B, D)`.
                                Data type is float.

        Returns:
            Tensor, has the same data type and shape as original `pbc_box`.
        """
        if self._pbc_box is None:
            return pbc_box
        self._pbc_box.copy_(pbc_box)
        return self._pbc_box

    def get_kinetics(self, velocity: Tensor) -> Tensor:
        r"""
        Calculate kinetics according to velocity.

        Args:
            velocity(Tensor):   Tensor of atomic velocities. Tensor shape is `(B, A, D)`.
                                Data type is float.

        Returns:
            Tensor, Tensor of kinetics. Tensor shape is `(B, D)`. Data type is float.
        """
        if velocity is None:
            return None
        # (B, A, D) * (B, A, 1)
        k = 0.5 * self._atom_mass * velocity**2
        # (B, D) <- (B, A, D)
        kinetics = torch.sum(k, dim=-2)
        return kinetics * self.kinetic_unit_scale

    def get_temperature(self, kinetics: Tensor = None) -> Tensor:
        r"""
        Calculate temperature according to velocity.

        Args:
            kinetics(Tensor):   Tensor of kinetics. Tensor shape is `(B, D)`. Data type is float. Default: None

        Returns:
            Tensor, Tensor of temperature. The shape of the Tensor is `(B)`. Data type is float.
        """
        if kinetics is None:
            return None
        # (B) <- (B, D)
        kinetics = torch.sum(kinetics, dim=-1)
        return 2 * kinetics / self.degrees_of_freedom / self.boltzmann

    def get_volume(self, pbc_box: Tensor) -> Tensor:
        r"""
        Calculate volume according to PBC box

        Args:
            pbc_box(Tensor):    Tensor of PBC box. Tensor shape is `(B, D)`. Data type is float.

        Returns:
            Tensor, Tensor of volume. The shape of the Tensor is `(B)`, and the data type is float.
        """
        if self._pbc_box is None:
            return None
        # (B, 1) <- (B, D)
        return torch.prod(pbc_box, -1, keepdim=True)

    def get_pressure(self, kinetics: Tensor, virial: Tensor, pbc_box: Tensor) -> Tensor:
        r"""
        Calculate pressure according to kinetics, viral and PBC box.

        Args:
            kinetics(Tensor):   Tensor of kinetics. Tensor shape is `(B, D)`. Data type is float.
            virial(Tensor):     Tensor of virial. Tensor shape is `(B, D)`. Data type is float.
            pbc_box(Tensor):    Tensor of PBC box. Tensor shape is `(B, D)`. Data type is float.

        Returns:
            Tensor, Tensor of pressure. Tensor shape is `(B, D)`. Data type is float.
        """
        if self._pbc_box is None:
            return None
        volume = torch.prod(pbc_box, -1, keepdim=True)
        # (B, D) = ((B, D) - (B, D)) / (B, 1)
        pressure = 2 * (kinetics - virial) / volume
        return pressure * self.press_unit_scale

    def get_com(self, coordinate: Tensor, keepdims: bool = True) -> Tensor:
        r"""
        Get coordinate of center of mass.

        Args:
            coordinate(Tensor): Tensor of atomic coordinates. Tensor shape is `(B, A, D)`.
                                Data type is float.
            keepdims(bool):    Whether to keep the last dimension. Default: True

        Returns:
            Tensor, Tensor of center of mass. Tensor shape is `(B, D)` or `(B, 1, D)`.
                Data type is float.
        """
        # (B, A, D) * (B, A, 1)
        com = coordinate * self._atom_mass
        # (B, D) or (B, 1, D) <- (B, A, D)
        com = torch.sum(com, dim=-2, keepdim=keepdims)
        # (B, D) or (B, 1, D) = (B, D) or (B, 1, D) / (B, 1)
        if keepdims:
            com = com / self.system_natom.unsqueeze(-1)
        else:
            com = com / self.system_natom
        return com

    def get_com_velocity(self, velocity: Tensor, keepdims: bool = True) -> Tensor:
        r"""
        Get velocity of center of mass.

        Args:
            velocity(Tensor):   Tensor of atomic velocities. Tensor shape is `(B, A, D)`.
                                Data type is float.
            keepdims(bool):    Whether to keep the last dimension. Default: True

        Returns:
            Tensor, Tensor of center of mass velocity. Tensor shape is `(B, D)` or `(B, 1, D)`.
                Data type is float.
        """
        # (B, A, D) * (B, A, 1)
        com_vel = velocity * self._atom_mass
        # (B, D) or (B, 1, D) <- (B, A, D)
        com_vel = torch.sum(com_vel, dim=-2, keepdim=keepdims)
        # (B, D) or (B, 1, D) = (B, D) or (B, 1, D) / (B, 1)
        if keepdims:
            com_vel = com_vel / self.system_mass.unsqueeze(-1)
        else:
            com_vel = com_vel / self.system_mass
        return com_vel

    def forward(self,
               coordinate: Tensor,
               velocity: Tensor,
               force: Tensor,
               energy: Tensor,
               virial: Tensor = None,
               pbc_box: Tensor = None,
               steps: int = 0,
               **kwargs
               ) -> Dict[str, Tensor]:
        r"""
        Forward function of controller.

        Args:
            coordinate(Tensor): Tensor of atomic coordinates. Tensor shape is `(B, A, D)`.
                                Data type is float.
            velocity(Tensor):   Tensor of atomic velocities. Tensor shape is `(B, A, D)`.
                                Data type is float.
            force(Tensor):      Tensor of atomic forces. Tensor shape is `(B, A, D)`.
                                Data type is float.
            energy(Tensor):     Tensor of system energy. Tensor shape is `(B, 1)`.
                                Data type is float.
            virial(Tensor):     Tensor of virial. Tensor shape is `(B, D)`. Data type is float.
                                Default: None
            pbc_box(Tensor):    Tensor of PBC box. Tensor shape is `(B, D)`. Data type is float.
                                Default: None
            step(int):          Simulation step. Default: 0

        Returns:
            Dict[str, Tensor], Dictionary of output tensors.
        """
        raise {'coordinate': coordinate,
                'velocity': velocity,
                'force': force,
                'energy': energy,
                'virial': virial,
                'pbc_box': pbc_box,
                }

    def _get_mw_tensor(self,
                      value: Union[float, ndarray, Tensor, List[float]],
                      name: str) -> Tensor:
        """Get tensor for multiple walkers"""

        value = get_tensor(value, dtype=torch.float32)

        if value.ndim < 1:
            return value.reshape(-1)
        elif value.shape[0] == 1:
            return value.reshape(-1)

        if value.shape[0] != self.num_walker:
            error_info = f'The size of {name} must be equal to 1 or ' \
                         f'the number of multiple walker ({self.num_walker}) but got '
            raise ValueError(error_info)

        return value.reshape(self.num_walker, 1)

    def print_info(self) -> Dict:
        """Print controller information and return as a dictionary.

        Returns:
            dict: A dictionary containing controller information
        """
        info = {
            'type': self.__class__.__name__,
            'control_step': int(self.control_step),
            'time_step': float(self.time_step),
            'degrees_of_freedom': int(self.degrees_of_freedom),
            'num_constraints': int(self.num_constraints),
        }
        
        return info
    
